#!/usr/bin/python
import json
import os
import sys

sys.setrecursionlimit(4000)
# json_txt = open('D:/CBG/DFX/test/memleak-js-com.example.testdemo-9464-20240118114338.heapsnapshot', 'r').read()
# json_txt = open('D:/CBG/DFX/test/memleak-js-com.example.testdemo-13255-20240130111436.heapsnapshot', 'r').read()
# json_txt = open('D:/CBG/DFX/test/closure.heapsnapshot', 'r', encoding='utf-8-sig').read()
json_txt = open('D:/CBG/DFX/test/tmp_test2.heapsnapshot', 'r', encoding='utf-8-sig').read()
json_obj = json.loads(json_txt)

node_fields_checker = ["type","name","id","self_size","edge_count","trace_node_id","detachedness"]


def list_to_str(list_data, separator):
    return separator.join(str(item) for item in list_data)


def list_to_str_with_flag(list_data, separator, flag):
    return separator.join(str(flag) + "_" + str(item) for item in list_data)


def write_str_to_file(file_name, write_text_str, mod='w'):
    file_dir = os.path.dirname(file_name)
    if not os.path.exists(file_dir):
        os.makedirs(file_dir)
    op_file = open(file_name, mod, encoding='utf-8')
    op_file.writelines(write_text_str)
    op_file.close()


class DrawMdGraph:
    """
    direct: TB	从上到下
            BT	从下到上
            LR	从左到右
            RL	从右到左
    """
    def __init__(self, direct="TB"):
        self.block_start = '```mermaid\n'
        self.block_end = '\n```'
        self.graph_type = "graph "
        self.direct = direct
        self.graph_node_str_list = []

    def get_graph_str(self):
        ret = self.block_start + self.graph_type + self.direct + "\n"
        ret += list_to_str(self.graph_node_str_list, '\n')
        ret += self.block_end
        return ret

    def add_directed_graph_node(self, f_id, t_id, f_str="", t_str="", edge_txt=""):
        from_nd = str(f_id)
        if f_str is not None and f_str != "":
            from_nd += "[" + f_str + "]"
        edge_str = "--"
        if edge_txt is not None and edge_txt != "":
            edge_str += edge_txt + "-->"
        else:
            edge_str += ">"
        to_nd = str(t_id)
        if t_str is not None and t_str != "":
            to_nd += "[" + t_str + "]"
        ret = from_nd + edge_str + to_nd
        self.graph_node_str_list.append(ret)

    def add_directed_graph_node_show_id(self, f_id, t_id, f_str="", t_str="", edge_txt=""):
        from_nd = str(f_id)
        if f_str is not None and f_str != "":
            from_nd += "[" + str(f_id) + ": " + f_str + "]"
        edge_str = "--"
        if edge_txt is not None and edge_txt != "":
            edge_str += edge_txt + "-->"
        else:
            edge_str += ">"
        to_nd = str(t_id)
        if t_str is not None and t_str != "":
            to_nd += "[" +  str(t_id) + ": " + t_str + "]"
        ret = from_nd + edge_str + to_nd
        self.graph_node_str_list.append(ret)


class HeapOneNode:
    def __init__(self, node_fields, node_data_list, node_types_list, strings_list):
        self.node_fields_list = node_fields.copy() # ["type","name","id","self_size","edge_count","trace_node_id","detachedness","native_size"]
        self.node_data_list = node_data_list.copy() # [2,3,3,26,0,0,0,0]
        self.field_type = node_data_list[0]
        self.field_name = node_data_list[1]
        self.field_id = node_data_list[2]
        self.field_self_size = node_data_list[3]
        self.field_edge_count = node_data_list[4]
        self.field_trace_node_id = node_data_list[5]
        self.field_detachedness = node_data_list[6]
        self.str_name = "NOT_SET"
        self.str_type = "NOT_SET"
        if len(node_data_list) <= 7:
            self.field_native_size = -1
        else:
            self.field_native_size = node_data_list[7]
        if self.field_type > len(node_types_list[0]):
            print("Err: 节点type错误 id=%s, type=%s"% (self.field_id, self.field_type))
        else:
            self.str_type = node_types_list[0][self.field_type]
        if self.str_type == "number":
            self.str_name = self.field_name
        else:
            tmp_ind = strings_list[self.field_name].find("\n")
            if tmp_ind != -1:
                self.str_name = strings_list[self.field_name][:tmp_ind:]
            else:
                self.str_name = strings_list[self.field_name]
            self.str_name = self.str_name.replace("\r", " ")
            self.str_name = self.str_name.replace("\t", " ")
            if self.str_name == '"':
                self.str_name = self.str_name.replace('"', "``")
            # if self.field_id == 11661:
            #     tmp = strings_list[self.field_name]
            #     print(self.field_id)

        self.node_fields_list.append("name_type")
        self.node_data_list.append(self.str_type)
        self.node_fields_list.append("str_name")
        self.node_data_list.append(self.str_name)
        self.root_dist = -1
        self.root_reachable = False
        self.node_fields_list.append("dist")
        self.node_fields_list.append("reach")
        self.node_data_list.append(self.root_dist)
        self.node_data_list.append("N")
        self.field_cnt_for_dump_edge = len(self.node_fields_list) # 与边汇合时，用于截断，不打印node cnt
        # 用于构建节点间关系 from_node --from_edge--> curr_node(self) --to_edge--> to_edge
        self.heap_from_node_obj_list = []
        self.heap_from_edge_obj_list = []
        self.heap_to_edge_obj_list = []
        self.heap_to_node_obj_list = []
        self.node_fields_list.append("from_node_cnt")
        self.node_data_list.append(0) # 初始值为0
        self.node_fields_list.append("to_node_cnt")
        self.node_data_list.append(0) # 初始值为0
        self.have_searched = False

    def set_dist_to_root(self, dist, is_from_root):
        if self.root_reachable and not is_from_root:
            print("%s,%s 已设可达root = %d，收到一不可达设置[%d, %s]\n" % (self.field_id, self.str_name, self.root_dist, dist, is_from_root))
            return
        if self.root_dist != -1 and dist >= self.root_dist:
            return
        self.root_dist = dist
        self.root_reachable = is_from_root
        self.node_data_list[-4] = self.root_dist
        if is_from_root:
            self.node_data_list[-3] = "Y"

    def set_from_node(self, heap_one_edge_obj, heap_one_node_obj):
        self.heap_from_edge_obj_list.append(heap_one_edge_obj)
        self.heap_from_node_obj_list.append(heap_one_node_obj)
        # 每添加一个from节点，对应from_node_cnt +1
        self.node_data_list[-2] += 1

    def set_to_node(self, heap_one_edge_obj, heap_one_node_obj):
        self.heap_to_edge_obj_list.append(heap_one_edge_obj)
        self.heap_to_node_obj_list.append(heap_one_node_obj)
        self.node_data_list[-1] += 1

    def to_str(self):
        ret = ""
        for i in range(len(self.node_fields)):
            ret += self.node_fields[i] + ": " + self.node_data_list[i] + ", "
        return ret

    def __str__(self):
        return self.to_str()

    def get_node_fields_in_str(self, separator):
        return list_to_str(self.node_fields_list, separator)

    def get_node_contents_in_str(self, separator):
        return list_to_str(self.node_data_list, separator)

    def get_node_fields_str_for_dump_edge(self, separator, flag="F"):  # F: from T: to
        tmp_list = self.node_fields_list[:self.field_cnt_for_dump_edge:]
        return list_to_str_with_flag(tmp_list, separator, flag)

    def get_node_contents_str_for_dump_edge(self, separator):
        tmp_list = self.node_data_list[:self.field_cnt_for_dump_edge:]
        return list_to_str(tmp_list, separator)


class HeapNodeList:
    """
    json_node_fields:"node_fields"(不含):["type","name","id","self_size","edge_count","trace_node_id","detachedness","native_size"]
    json_nodes: "nodes"（不含）:[9,15181,1,0,4816,0,0,0...]
    """
    def __init__(self, json_node_fields, json_node_types, json_strings, node_count, json_nodes):
        self.node_fields_list = json_node_fields #["type","name","id","self_size","edge_count","trace_node_id","detachedness","native_size"]
        self.node_types_list = json_node_types
        self.strings_list = json_strings
        self.node_count = node_count
        self.json_nodes = json_nodes
        self.nodes_list = [] # 保留分行后的原始数据，[[1,2,3,4,5,6],[1,2,3,4,5,6]]
        self.heap_one_node_obj_map = {} # 保存构建成HeapOneNode对象的map {id: obj}
        if node_count != len(json_nodes) / len(self.node_fields_list):
            print("Err: 节点数和节点量不匹配，%d, %d / %d = %d\n" %
                  (node_count, len(json_nodes), len(self.node_fields_list), len(json_nodes) / len(self.node_fields_list)))
        self.root_node = None
        for i in range(node_count):
            start_ind = i * len(self.node_fields_list)
            end_ind = start_ind + len(self.node_fields_list)
            tmp_one_node_data = json_nodes[start_ind:end_ind:]
            # print(i, start_ind, end_ind, tmp_one_node_data)
            self.nodes_list.append(tmp_one_node_data)
            heap_one_node_tmp = HeapOneNode(json_node_fields, tmp_one_node_data, self.node_types_list, self.strings_list)
            self.heap_one_node_obj_map[heap_one_node_tmp.field_id] = heap_one_node_tmp
            if self.root_node is None and (heap_one_node_tmp.str_name == "SyntheticRoot" or heap_one_node_tmp.field_id == 1):
                self.root_node = heap_one_node_tmp

    def node_fields_to_str(self, separator):
        return list_to_str(self.node_fields_list, separator)

    def to_text_str(self, separator):
        ret = self.node_fields_to_str(separator) + "\n"
        for i in range(self.node_count):
            ret += list_to_str(self.nodes_list[i], separator) + "\n"
        return ret

    def __str__(self):
        return self.to_text_str(',')

    def get_heap_one_node_obj_fields_str(self, separator):
        if len(self.heap_one_node_obj_map.keys()) == 0:
            return self.node_fields_to_str(separator) + "\n"
        id = list(self.heap_one_node_obj_map.keys())[0]
        return self.heap_one_node_obj_map[id].get_node_fields_in_str(separator) + "\n"

    def get_heap_one_node_obj_list_in_str(self, separator):
        ret = ""
        for i in self.heap_one_node_obj_map.keys():
            ret += self.heap_one_node_obj_map[i].get_node_contents_in_str(separator) + "\n"
        return ret

    def get_all_heap_one_node_obj_in_str(self, separator):
        return self.get_heap_one_node_obj_fields_str(separator) + self.get_heap_one_node_obj_list_in_str(separator)


class HeapOneEdge:
    def __init__(self, json_edge_fields, edges_data_list, json_edge_types, strings):
        self.edge_data_list = edges_data_list.copy()
        self.edge_fields = json_edge_fields.copy()
        self.field_type = edges_data_list[0]
        self.field_name_or_index = edges_data_list[1]
        self.field_to_node = edges_data_list[2]
        self.type_name = json_edge_types[0][int(self.field_type)]
        if self.field_type == "element":
            self.str_name = self.field_name_or_index
        else:
            self.str_name = strings[int(self.field_name_or_index)]
        self.edge_data_list.append(self.type_name)
        self.edge_data_list.append(self.str_name)
        self.edge_fields.append("type_name")
        self.edge_fields.append("str_name")
        self.heap_one_node_obj_from = None
        self.heap_one_node_obj_to = None

    def get_edge_fields_in_str(self, separator):
        node_from_str = self.heap_one_node_obj_from.get_node_fields_str_for_dump_edge(separator, "F")
        node_to_str = "To" + separator + self.heap_one_node_obj_from.get_node_fields_str_for_dump_edge(separator, "T")
        return node_from_str + separator + list_to_str(self.edge_fields, separator) + separator + node_to_str

    def get_edge_data_in_str(self, separator):
        node_from_str = self.heap_one_node_obj_from.get_node_contents_str_for_dump_edge(separator)
        if self.heap_one_node_obj_to is None:
            node_to_str = "N"
        else:
            node_to_str = "Y" + separator + self.heap_one_node_obj_to.get_node_contents_str_for_dump_edge(separator)
        return node_from_str + separator + list_to_str(self.edge_data_list, separator) + separator + node_to_str

    def set_from_and_to_node_obj(self, heap_one_node_obj_from, heap_one_node_obj_to):
        if heap_one_node_obj_from is None:
            print("Err: 边的from节点不可能是None, 边信息：%s\n" % self.get_edge_data_in_str(","))
        self.heap_one_node_obj_from = heap_one_node_obj_from
        self.heap_one_node_obj_to = heap_one_node_obj_to

    def gen_markdown_fmt_str(self, markdown_graph_obj):
        # markdown_graph_obj = DrawMdGraph()
        f_nd = self.heap_one_node_obj_from
        t_nd = self.heap_one_node_obj_to
        markdown_graph_obj.add_directed_graph_node_show_id(f_nd.field_id, t_nd.field_id, f_nd.str_name, t_nd.str_name, self.str_name)


class HeapEdgeList:
    """
    json_node_fields:"edge_fields"(不含):["type","name_or_index","to_node"]
    json_nodes: "edges"（不含）:[5,15182,358240...]
    """
    def __init__(self, json_edge_fields, json_edge_types, edge_count, json_edges, strings):
        self.edge_fields = json_edge_fields # ["type","name_or_index","to_node"]
        self.edge_types_list = json_edge_types # [["context","element","property","internal","hidden","shortcut","weak"],"string_or_number","node"]
        self.edge_count = edge_count
        self.edges_list = [] # 保留分行后的原始数据，[[1,2,3],[1,2,3]]
        self.heap_one_edge_obj_list = []
        for i in range(edge_count):
            start_ind = i * len(self.edge_types_list)
            end_ind = start_ind + len(self.edge_types_list)
            tmp_one_node_data = json_edges[start_ind:end_ind:]
            self.edges_list.append(tmp_one_node_data)
            tmp_heap_edge_obj = HeapOneEdge(json_edge_fields, tmp_one_node_data, json_edge_types, strings)
            self.heap_one_edge_obj_list.append(tmp_heap_edge_obj)

    def node_fields_to_str(self, separator):
        return list_to_str(self.edge_types_list, separator)

    def to_text_str(self, separator):
        ret = self.node_fields_to_str(separator) + "\n"
        for i in range(self.edge_count):
            ret += list_to_str(self.edges_list[i], separator) + "\n"
        return ret

    def __str__(self):
        return self.to_text_str(',')

    def get_all_heap_one_edge_obj_in_str(self, separator):
        if len(self.heap_one_edge_obj_list) == 0:
            ret = self.node_fields_to_str(separator) + "\n"
        else:
            ret = self.heap_one_edge_obj_list[0].get_edge_fields_in_str(separator) + "\n"
        for i in self.heap_one_edge_obj_list:
            ret += i.get_edge_data_in_str(separator) + "\n"
        return ret


def update_dist_to_root_node(heap_one_node_obj, curr_dist, is_from_root):
    # print("%d,%s,%s,%s" % (curr_dist, is_from_root, heap_one_node_obj.field_id, heap_one_node_obj.str_name))
    # heap_one_node_obj.set_dist_to_root(curr_dist, is_from_root)
    # if curr_dist > 1000:
    #     print("Err: 距离超过1000")
    #     return
    # for next_nd in heap_one_node_obj.heap_to_node_obj_list:
    #     if next_nd.root_dist != -1:
    #         continue
    #     update_dist_to_root_node(next_nd, curr_dist + 1, is_from_root)
    curr_list = [heap_one_node_obj]
    next_list = []
    while len(curr_list) != 0:
        for curr_nd in curr_list:
            # print("%d,%s,%s,%s" % (curr_dist, is_from_root, curr_nd.field_id, curr_nd.str_name))
            curr_nd.set_dist_to_root(curr_dist, is_from_root)
            for next_nd in curr_nd.heap_to_node_obj_list:
                if next_nd.root_dist != -1 or next_nd in next_list:
                    continue
                next_list.append(next_nd)
        curr_list = next_list
        next_list = []
        curr_dist += 1


def search_start_node(heap_one_node_obj):
    # if len(heap_one_node_obj.heap_from_node_obj_list) == 0: # 如果当前节点无from节点，说明是一独立起始节点
    #     update_dist_to_root_node(heap_one_node_obj, 0, False) # 更新该节点的所有节点
    #     return
    # # 如果不是起始节点，则继续找其上一层节点
    # for from_nd in heap_one_node_obj.heap_from_node_obj_list:
    #     search_start_node(from_nd)
    curr_list = [heap_one_node_obj]
    next_list = []
    level0_nodes = []
    if heap_one_node_obj.field_id == 88257 or heap_one_node_obj.field_id == 15073:
        print("unreach obj")
    while len(curr_list) != 0:
        # print("当前长度：",len(curr_list))
        for curr_nd in curr_list:
            if len(curr_nd.heap_from_node_obj_list) == 0:
                if curr_nd not in level0_nodes:
                    # print("增加一节点：", curr_nd.field_id, curr_nd.str_name)
                    level0_nodes.append(curr_nd)
                continue
            if curr_nd.root_reachable or curr_nd.have_searched:
                continue
            if curr_nd.field_id == 15073:
                print("unreach obj")
            curr_nd.have_searched = True
            for from_nd in curr_nd.heap_from_node_obj_list:
                if from_nd in next_list:
                    continue
                next_list.append(from_nd)
        if len(next_list) == 0:
            for curr_nd in curr_list:
                if curr_nd not in level0_nodes:
                    # print("增加一节点：", curr_nd.field_id, curr_nd.str_name)
                    level0_nodes.append(curr_nd)
        curr_list = next_list
        next_list = []
    for level0_nd in level0_nodes:
        # print("其他头节点：%s,%s" % (curr_nd.field_id, curr_nd.str_name))
        update_dist_to_root_node(level0_nd, 0, False)


node_output_file_path = "D:/CBG/DFX/test/00_nodes_info2.txt"
edge_output_file_path = "D:/CBG/DFX/test/01_edges_info2.txt"
def analysis_data(json_obj):
    snapshot_head = json_obj.get("snapshot")
    meta_map = snapshot_head.get("meta")
    node_fields = meta_map.get("node_fields")
    node_types = meta_map.get("node_types")
    nodes = json_obj.get("nodes")
    strings = json_obj.get("strings")
    node_count = int(snapshot_head.get("node_count"))
    heap_node_list = HeapNodeList(node_fields, node_types, strings, node_count, nodes)
    edges = json_obj.get("edges")
    edge_count = int(snapshot_head.get("edge_count"))
    edge_fields = meta_map.get("edge_fields")
    edge_types = meta_map.get("edge_types")
    heap_edge_list = HeapEdgeList(edge_fields, edge_types, edge_count, edges, strings)
    curr_edge_ind = 0
    for node_id in heap_node_list.heap_one_node_obj_map.keys():
        curr_heap_one_node_obj = heap_node_list.heap_one_node_obj_map.get(node_id)
        curr_node_edge_cnt = curr_heap_one_node_obj.field_edge_count
        for i in range(curr_node_edge_cnt):
            curr_edge_obj = heap_edge_list.heap_one_edge_obj_list[curr_edge_ind + i]
            to_node_id = heap_node_list.json_nodes[curr_edge_obj.field_to_node + 2]
            # print(to_node_id)
            to_node_obj = heap_node_list.heap_one_node_obj_map.get(to_node_id)
            curr_edge_obj.set_from_and_to_node_obj(curr_heap_one_node_obj, to_node_obj)
            if to_node_obj is None:
                continue
            curr_heap_one_node_obj.set_to_node(curr_edge_obj, to_node_obj)
            to_node_obj.set_from_node(curr_edge_obj, curr_heap_one_node_obj)
        curr_edge_ind += curr_node_edge_cnt
    print("所有节点建图完成")
    if curr_edge_ind != heap_edge_list.edge_count:
        print("Err: 累加边数量不等于记录边数量, 累加：%d, 原文件记录：%d\n" % (curr_edge_ind, heap_edge_list.edge_count))
    if heap_node_list.root_node is None:
        print("Err: 未找到根节点")
        return
    print("计算至根节点距离...")
    update_dist_to_root_node(heap_node_list.root_node, 0, True)
    print("遍历至根节点距离结束")
    print("寻找其他不可达头节点，并计算距离")
    for k_id in heap_node_list.heap_one_node_obj_map.keys():
        curr_node = heap_node_list.heap_one_node_obj_map.get(k_id)
        if not curr_node.root_reachable and curr_node.root_dist == -1:  # root不可达根节点
            # print("节点：", curr_node.field_id, curr_node.str_name)
            search_start_node(curr_node)
        if not curr_node.root_reachable and curr_node.root_dist == -1 and curr_node.str_name=="HiddenClass(NonMovable)":
            search_start_node(curr_node)
            print("hidden class is -1")
    print("其他不可达头节点计算距离结束")
    write_str_to_file(node_output_file_path, heap_node_list.get_all_heap_one_node_obj_in_str('\t'))
    write_str_to_file(edge_output_file_path, heap_edge_list.get_all_heap_one_edge_obj_in_str('\t'))
    draw_markdown_graph_obj = DrawMdGraph("LR")
    for curr_edge in heap_edge_list.heap_one_edge_obj_list:
        curr_edge.gen_markdown_fmt_str(draw_markdown_graph_obj)
    print(draw_markdown_graph_obj.get_graph_str())


if __name__ == '__main__':
    # print(json_obj)
    print(type(json_obj))
    analysis_data(json_obj)
